#pragma once

#include "ros/ros.h"
#include <tf2_ros/transform_broadcaster.h>
#include <geometry_msgs/TransformStamped.h>
#include <nav_msgs/Odometry.h>


class odomExtrapolation
{
private:
    ros::Subscriber odom_sub;
    tf2_ros::TransformBroadcaster br;
    nav_msgs::Odometry odom_new, odom_old;
    ros::WallTimer timer;
    ros::Publisher odom_pub;
    
    double vx, vy, vz;
    double rate;
    double dt;
    double timeshift;

    bool publish_odom, enable_time_sync;

    std::string frame_id;

public:
    odomExtrapolation(ros::NodeHandle *nh)
    {
        auto hints = ros::TransportHints().tcpNoDelay();
        std::string odom_topic;

        nh->param<std::string>("odom_topic", odom_topic, "/Odometry");
        nh->param<std::string>("frame_id", frame_id, "body_ext");
        nh->param<double>("rate", rate, 100.0);
        nh->param<bool>("publish_odom", publish_odom, true);
        nh->param<bool>("enable_time_sync", enable_time_sync, true);

        dt = 1.0 / rate;

        odom_sub = nh->subscribe(odom_topic, 10, &odomExtrapolation::odom_callback, this, hints);
        timer = nh->createWallTimer(ros::WallDuration(dt), &odomExtrapolation::extrapolation_callback, this);
        odom_pub = nh->advertise<nav_msgs::Odometry>("/odom_ext", 100);

        timer.stop();
    }

    void odom_callback(const nav_msgs::Odometry::ConstPtr &msg)
    {
        odom_old = odom_new;
        odom_new = *msg;

        double delta_t = odom_new.header.stamp.toSec() - odom_old.header.stamp.toSec();
        timeshift = enable_time_sync ? ros::Time::now().toSec() - odom_new.header.stamp.toSec() : 0.0;
        delta_t = delta_t > 0.001 ? delta_t : 0.1;
        vx = (odom_new.pose.pose.position.x - odom_old.pose.pose.position.x) / delta_t;
        vy = (odom_new.pose.pose.position.y - odom_old.pose.pose.position.y) / delta_t;
        vz = (odom_new.pose.pose.position.z - odom_old.pose.pose.position.z) / delta_t;
    }

    void extrapolation_callback(const ros::WallTimerEvent& event)
    {
        auto t = ros::Time::now();
        double delta_t = t.toSec() - odom_new.header.stamp.toSec() - timeshift;
        if (delta_t > 0.2)
            ROS_INFO_STREAM("dt: " << delta_t << ", vx: " << vx << ", timeshift: " << timeshift);

        geometry_msgs::TransformStamped transformStamped;
        transformStamped.header.stamp = t;
        transformStamped.header.frame_id = odom_new.header.frame_id;
        transformStamped.child_frame_id = frame_id;
        transformStamped.transform.translation.x = odom_new.pose.pose.position.x + delta_t * vx;
        transformStamped.transform.translation.y = odom_new.pose.pose.position.y + delta_t * vy;
        transformStamped.transform.translation.z = odom_new.pose.pose.position.z + delta_t * vz;
        transformStamped.transform.rotation.w = odom_new.pose.pose.orientation.w;
        transformStamped.transform.rotation.x = odom_new.pose.pose.orientation.x;
        transformStamped.transform.rotation.y = odom_new.pose.pose.orientation.y;
        transformStamped.transform.rotation.z = odom_new.pose.pose.orientation.z;

        br.sendTransform(transformStamped);

        if (publish_odom)
        {
            nav_msgs::Odometry msg;
            msg = odom_new;
            msg.header.stamp = t;
            msg.pose.pose.position.x += delta_t * vx;
            msg.pose.pose.position.y += delta_t * vy;
            msg.pose.pose.position.z += delta_t * vz;

            odom_pub.publish(msg);
        }
    }

    void start()
    {
        timer.start();
        ros::spin();
    }

};
